# TODO: 导入必要的库和模块
import gradio as gr
import pickle
import numpy as np
from PIL import Image
import os
import matplotlib.pyplot as plt

# 获取当前工作目录
# TODO: 加载保存的KNN模型，这样我们可以使用预训练的模型进行预测
best_knn = pickle.load(open('best_knn_model.pkl', 'rb'))
# TODO: 定义预测函数，这个函数将用于Gradio接口进行预测
def predict_digit(image):
    image = np.array(image['layers'][0])
    image = image[:, :, 3]
    image = Image.fromarray(image)
    image = image.resize((8,8))
    image = np.array(image)
    image = (image/ 255) * 16
    image = image.reshape(1, -1)
    prediction = best_knn.predict(image)
    return int(prediction[0])
# TODO: 创建Gradio接口，这个接口将用于用户输入和显示预测结果
sketchpad = gr.Sketchpad()
gr.Interface(fn=predict_digit, inputs=sketchpad, outputs=gr.Label()).launch()
# TODO: 启动Gradio接口，用户可以通过这个接口进行交互
